# -*- coding:utf-8 -*-
# @Time: 2021/1/10 20:31
# @Author: 周铖鑫
# @Description：pytorch概述
# @File: UsePytorch.py
import torch
import numpy as np
import  pandas as pd
import matplotlib.pyplot as plt
from torch import nn

plt.rcParams['font.family'] = ['Microsoft YaHei']
plt.rcParams['font.sans-serif'] = ['SimHei']
data=pd.read_csv("income.csv")
plt.scatter(data.Education,data.Income)
plt.xlabel("教育")
plt.ylabel("收入")
# print(data)
data.info()


X=torch.from_numpy(data.Education.values.reshape(-1,1).astype(np.float32))
Y=torch.from_numpy(data.Income.values.reshape(-1,1).astype(np.float32))
model=nn.Linear(1,1)#out=w*input+b等价model(input)
loss_fn=nn.MSELoss()
optimizer=torch.optim.SGD(model.parameters(),lr=0.0001)

'''training'''
for epoch in range(5000):
    for x,y in zip(X,Y):
        y_pred=model(x) #使用模型预测
        loss=loss_fn(y,y_pred) #根据预测结果计算损失
        optimizer.zero_grad() #把变量梯度变为0
        loss.backward() #求解梯度，反向传播算法
        optimizer.step() #优化模型参数

#保存
torch.save(model, 'model.pkl')
# 加载
# model = torch.load('\model.pkl')

 # 保存模型参数
torch.save(model.state_dict(), 'parameter.pkl')
 # 加载
# model = TheModelClass(...)
# model.load_state_dict(torch.load('parameter.pkl'))
print('weight:{}    bias:{}'.format(model.weight, model.bias))
plt.plot(X.numpy(),model(X).data.numpy(),c="r")
plt.show()